import sys
import asyncio

from flask import Flask, request, jsonify

from model.glm_model import GLMModel
from model.openai_model import OpenAIModel
from server.server_api_param_parser import ServiceArgumentParser
from translator.pdf_translator import PDFTranslator
from utils.config_loader import ConfigLoader

app = Flask(__name__)


@app.route('/translate_pdf', methods=['POST'])
def translate_pdf():
    # 解析参数
    parser = ServiceArgumentParser()
    args = parser.parse_args()

    # 获取上传的PDF文件
    file = request.files['pdf_file']
    file.save('input.pdf')

    # 执行翻译操作，生成新的PDF文件
    output_file_path = translate_pdf_function(args, "./input.pdf")

    # 返回新生成的PDF文件地址
    return jsonify({'translated_pdf_path': output_file_path})


def translate_pdf_function(args, input_pdf_path):
    config_loader = ConfigLoader("../config.yaml")
    config = config_loader.load_config()
    model_name = args.openai_model if args.openai_model else config['OpenAIModel']['model']
    if model_name == "OpenAIModel":
        api_key = args.openai_api_key if args.openai_api_key else config['OpenAIModel']['api_key']
        model = OpenAIModel(model=model_name, api_key=api_key)
    else:
        model_url = args.glm_model_url if args.glm_model_url else config['GLMModel']['model_url']
        model = GLMModel(model_url=model_url)
    translator = PDFTranslator(model)

    page_number = args.page_number if args.openai_api_key else sys.maxsize
    target_language = args.target_language if args.target_language else "中文"
    file_format = args.file_format if args.file_format else "PDF"

    # 获取上传的PDF文件
    file = request.files['pdf_file']
    file.save('input.pdf')

    translator.translate_pdf('./input.pdf', file_format, target_language, pages=page_number)
    output_file_path = translator.book.pdf_file_path.replace('.pdf', f'_translated.pdf')
    return output_file_path


if __name__ == '__main__':
    app.run()
